#include "itpg/itp.h"
#include "cktsat/cktcover.h"
#include "util/util.h"
#include "info/cktinfo.h"
#include "itpg/Solver.h"
#include "cktsat/satsolver.h"
#include <ctime>

#include <vector>
using std::vector;

typedef MinisatProof::vec< MinisatProof::Lit > Clause; // define the clause type
typedef MinisatProof::Lit plit_t;                      // define the literal type

/// the interface to dump the clause 
template<class Solver , class Clause>
static void addClause( Solver & s ,Clause & c , uint8_t p )
{ 
/*   for( int i = 0 ; i < c.size() ; ++i  )
      cout << (sign(c[i])?'-':' ') << var(c[i]) << " ";
   cout << "["<< (unsigned)p <<"]" << endl;*/
   s.addClause( c , p ) ; 
}

class ClauseSaver : public SATSolver
{
   vec< Clause > & c;
   public:
   ClauseSaver( vec<Clause > & in ):c(in){}
   virtual ~ClauseSaver(){}
   virtual void addClause( vec<Lit> & x)
   {
      c.push();
      Clause & y = c.last();
      y.growTo( x.size() );
      for( unsigned i = 0 ; i < x.size() ; ++i ) //TODO using copy here
         y[i] = x[i];
   }
};
/// cost estimation 
static double balance( unsigned np , unsigned npi , uint8_t ** p )
{
   // foreach part  cost = sum 1/#input
   double cost = 1;
   unsigned pi_num[ np ]; 
   memset( pi_num , 0 , sizeof( unsigned ) * np );

   for( unsigned i = 0 ; i < np; ++ i )
      for( unsigned j = 1 ; j < npi ; ++ j )
         if( p[j][i] ) 
            pi_num[i] += 1 ;

   std::sort( pi_num , pi_num + np );

   for( unsigned i = 0 ; i < np-1 ; ++ i )
      cost += pi_num[i+1] - pi_num[i];

   return 1.0/cost;
}
static double independent( unsigned np , unsigned npi , uint8_t ** p )
{
   // foreach input cost = sum 1/#part
   double cost = 0;
   for( unsigned j = 1 ; j < npi ; ++j )
   {
      unsigned n =1;
      for( unsigned i = 0 ; i < npi ; ++ i )
         if( p[j][i] ) 
            n+=1;
      cost += (1.0)/n;
   }
   return cost;
}
// utilization function for literal 
static plit_t shift( plit_t x , unsigned sh )
{
   return plit_t( var(x) + sh , sign(x) );
}
static void shift( const Clause & in , unsigned sh , Clause & out )
{
   out.clear(); out.growTo( in.size() );
   for( int i = 0 ; i < in.size() ; ++ i )
      out[i] = shift( in[i] , sh );
}


class Decomposition
{
   public:
   MinisatProof::Solver      _solver;
   CktCover                  _ckt ; 
   StrashAIG                 _itp;
   MinisatProof::MultiSITP * _mitp;
   
   // declare the function
   uint8_t          _fvar; ///< range from [1-16], restriction at onset and offset 
   vector<unsigned> _onset;
   vector<unsigned> _offset;
   
   uint8_t   **     _input; ///< input camp 
   plit_t    **     _part;  ///< assumption i , j means index i inside var j  assumption
   unsigned         _nvar;  ///< for assumpVar usage

   void assumpVar( MinisatProof::Lit l , unsigned & pinum , unsigned & part )const;
   void connect(  unsigned idx , vec<unsigned> & p );

   public:

   Decomposition(): _input(NULL) , _part(NULL) {};

   // initialize 
   void init( const char * );
   void read_fun( unsigned n , vector<unsigned> & o , vector<unsigned> & f );
   void dumpSolver();

   bool decompose();

   // report
   double cost()const{ return independent( _fvar , _ckt._andFlag , _input ) * balance( _fvar , _ckt._andFlag , _input ) ; }
   void report()const; 
   void model()const;  //sat assignment
   void conflict()const;
   void itp()const;
};
void Decomposition::init( const char * f )
{
   _ckt.read_aig( f );
   unsigned andFlag = _ckt._andFlag; 

   for( unsigned i = 1 ; i < andFlag ; ++ i )
      _itp.createBase();
}
// number of input , and the function 
void Decomposition::read_fun( unsigned n , vector<unsigned> & o ,  vector<unsigned> & f )
{
   unsigned andFlag = _ckt._andFlag; 

   _fvar = n;
   std::swap( o , _onset );
   std::swap( f , _offset );

   _part = new plit_t * [ andFlag ];
   for( unsigned i = 0 ; i < andFlag ; ++ i )
      _part[i] = new plit_t[n];
   _input = new uint8_t * [ andFlag ];
   for( unsigned i = 0 ; i < andFlag ; ++ i )
      _input[i] = new uint8_t[n];
   for( unsigned i = 0 ; i < andFlag ; ++i )
      for( unsigned j = 0 ; j < n ; ++ j )
         _input[i][j] =1;
}
bool Decomposition::decompose()
{
   Clause assumption;
   for( unsigned i = 1 ; i < _ckt._andFlag ; ++ i )
      for( int j = 0 ; j < _fvar ; ++ j )
         if( _input[i][j] )
            assumption.push( ~ _part[i][j] );
         else
            assumption.push( _part[i][j] ); // optional
   
/*   for( unsigned i = 0 ; i < assumption.size() ; ++i )
      cout << (sign(assumption[i])?'-':' ') << var(assumption[i]) << " ";
   cout << endl;*/

   bool ret = _solver.solve( assumption );
   if( ret == false )_mitp->done();
   return ret == false ;
}
// idx is one fvar , vars is shift value , p is connected part 
void Decomposition::connect(  unsigned idx  , vec<unsigned> & p )
{ 
   assert( p.size() > 1 );
   Clause c(3);
   for( unsigned i = 0 ; i < p.size() -1 ; ++ i)
      for( unsigned j = 1 ; j < _ckt._andFlag ; ++ j)
      {
         c[0] = _part[j][idx];
         Lit l = _ckt._lits[j];
         c[1] = shift( l , p[i]*_nvar);
         c[2] = ~shift( l , p[i+1]*_nvar);
         //_solver.addClause( c , 0xFF );
         addClause( _solver , c , 0xFF );
         c[1] = ~c[1] ;
         c[2] = ~c[2] ;
         //_solver.addClause( c , 0xFF );
         addClause( _solver , c , 0xFF );
      }
}
void Decomposition::dumpSolver()
{
   vec< Clause >  cls;
   ClauseSaver cs( cls );

   _ckt.init(0);     // no phase reducing 
   _ckt.covering(); 
   _ckt.assignLit( 0 );
   _ckt.assignCNF( cs );

   _nvar = _ckt.nVars();                                  // number of varialbe in each frame
   const unsigned andFlag = _ckt._andFlag;
   const unsigned nc = _onset.size() + _offset.size();    // number of circuit duplication  
   const unsigned nv = _nvar * nc + ( andFlag-1 ) * _fvar; //  
   
   { /// create Proof 
      _mitp = new MinisatProof::MultiSITP( nc );
      for( unsigned i = 0 ; i < nc ; ++ i )
      {
         unsigned sh = i*_nvar;

         MinisatProof::SMITP * p = new MinisatProof::SMITP( _itp , 0 );
         p->_camp = i;
         
         p->_base.growTo( nv );
         
         for( unsigned j = 0 ; j < nv ; ++ j ) 
            p->_base[j] = ((j/_nvar)==i)?MinisatProof::ITP::VarA : MinisatProof::ITP::VarB;
         for( unsigned j = 1 ; j < andFlag ; ++ j )
            p->_base[ MinisatProof::var(_ckt._lits[j]) + sh ] = toAigLit( j , MinisatProof::sign( _ckt._lits[j] ) ) ;

         _mitp->set( i , p );
      }
      _solver.proof = new MinisatProof::Proof( * _mitp );
   }

   { /// variable assignment 
      for( unsigned i = 0 ; i < nv ; ++ i )
         _solver.newVar();
      
      unsigned assumpt = _nvar * nc;

      for( int p = 0 ; p < _fvar ; ++ p )
         for( unsigned i = 1 ; i < andFlag ; ++ i )
            _part[i][p] = plit_t( assumpt ++ , 0 );
      assert( assumpt == nv );
   }

   { /// dump clause 
      Clause c;
      const plit_t f = _ckt._lits[ var( _ckt._po[0] ) ] ^ sign( _ckt._po[0] );

      for( unsigned i = 0 ; i < nc ; ++ i )
      {
         unsigned sh = i* _nvar;
         for( unsigned j = 0 ; j < cls.size() ; ++ j )
         {
            shift( cls[j] , sh ,  c );
            //_solver.addClause( c , i );
            addClause( _solver , c , i );
         }
         Clause v; v.push( (i < _onset.size()) ? ~shift( f , sh ) : shift(f,sh) );
         //_solver.addClause( v , i ); // i < onset , assign 0
         addClause( _solver , v , i );
      }

      for( int i = 0 , pos=1,neg=2 ; i < _fvar ; ++ i , pos=pos<<2, neg=neg<<2 )
      {
         vec<unsigned> p;

         for( unsigned j = 0 ; j < _onset.size() ; ++ j )
            if( _onset[j] & pos ) p.push( j );
         for( unsigned j = 0 ; j < _offset.size() ; ++ j)
            if( _offset[j] & neg )p.push( j + _onset.size() );

/*         if( p.size() )
            connect( i , vars , p );

         p.clear();*/
         for( unsigned j = 0 ; j < _onset.size() ; ++ j )
            if( _onset[j] & neg ) p.push( j );
         for( unsigned j = 0 ; j < _offset.size() ; ++ j)
            if( _offset[j] & pos )p.push( j+_onset.size() );
         
         if( p.size() )
            connect( i , p );
      }
   }
}
inline unsigned term( unsigned v , unsigned s ){ return 1<<(v+v+s) ;}
static void ordecomposition( Decomposition & d )
{
   vector<unsigned> o , f; 
   o.push_back( term(0,0) ); // a 
   o.push_back( term(1,0) ); // b
   f.push_back( term(0,1) | term(1,1) );// ~a~b

   d.read_fun( 2 , o , f );
}
static void andordecomposition( Decomposition & d )
{
   vector<unsigned> o , f ;
   o.push_back( term(0,0) | term(1,0) );
   o.push_back( term(2,0) );

   f.push_back( term(0,1) | term(2,1) );
   f.push_back( term(1,1) | term(2,1) );

   d.read_fun( 3 , o , f );
}
static void andorcomposition( Decomposition & d )
{

}
static double drand(){ return 1.0*(rand()%1000)/1000;}
static std::pair<unsigned,unsigned> permute( unsigned np , unsigned npi , uint8_t ** input )
{
   unsigned picnt[ npi ];
   unsigned pcnt[ np ];

   memset( picnt , 0 , sizeof(unsigned) * npi );
   memset( pcnt , 0 , sizeof(unsigned) * np );

   for( unsigned i = 1 ; i < npi ; ++ i )
      for( unsigned j = 0 ; j < np ; ++ j )
         if( input[i][j] )
            picnt[i]+=1 , pcnt[j]+=1;

   while(1)
   {
      unsigned a = rand()%(npi-1)+1;
      unsigned b = rand()%np;
      
      if( drand() < ( 1.0*picnt[a]*pcnt[b] / (np*npi) ) )
         return std::make_pair( a , b );
   }
}
static void reset( unsigned np , unsigned npi , uint8_t ** input )
{
   for( unsigned i = 1 ; i < npi ; ++ i)
      for( unsigned j = 0 ; j < np ; ++j )
         input[i][j] =0;
}
int main( unsigned argc , const char ** argv )
{
   ArgAnalyzer arg( "run <in:aig> [-seed <seed:int>]" ,  argc , argv );
   
   unsigned seed = time(0);
   if( arg["-seed"].setted ) seed = atoi( arg["-seed"][0] );
   cout << "seed " << seed << endl;
   srand( seed );

   Decomposition dcomp;
   dcomp.init( arg[0].c_str() );
   
   //ordecomposition( dcomp );
   andordecomposition( dcomp );

   dcomp.dumpSolver(); 

   for( unsigned i = 0 ; i < 100 ; ++ i )
   {
      std::pair<unsigned,unsigned> a = permute( dcomp._fvar , dcomp._ckt._andFlag , dcomp._input );

      dcomp._input[ a.first ][ a.second ] = 0;

      if( dcomp.decompose() == false ) 
         dcomp._input[ a.first ][ a.second ] = 1;
   }
   dcomp.report();
}
void Decomposition::model()const
{
   unsigned nc = _onset.size() + _offset.size() ;

   for( unsigned p = 0 ; p < nc ; ++ p )
   {
      unsigned shift = p * _ckt.nVars();
      for( unsigned i = 1 ; i < _ckt._andFlag ; ++ i )
         cout<< (_solver.model[ MinisatProof::var( _ckt._lits[i] ) + shift ] == MinisatProof::l_True) << " ";
      cout<<endl;
   }
}
void Decomposition::conflict()const
{
   
}
void Decomposition::assumpVar( MinisatProof::Lit l , unsigned & pinum , unsigned & part )const
{
   unsigned n = MinisatProof::var(l) - ( _nvar *(_onset.size()+_offset.size()) ) ;
   unsigned pis = _ckt._andFlag -1;

   pinum = (n%pis)+1;
   part  = n/pis;
}

void Decomposition::report()const
{
   cout << "independent " << independent( _fvar , _ckt._andFlag , _input ) <<endl;
   cout << "balance     " << balance( _fvar , _ckt._andFlag , _input ) <<endl;
   cout << "cost        " << cost() << endl;

   for( unsigned i = 0 ; i < _fvar ; ++ i )
   {
      for( unsigned j = 1 ; j < _ckt._andFlag ; ++j )
         cout << (unsigned)_input[j][i] ;
      cout<<endl;
   } 
}
void Decomposition::itp()const
{
   for( unsigned i = 0 ; i < _onset.size() ; ++i )
      cout << _mitp->_sitp[i]->interpolant() << " ";
   cout << " > ";
   for( unsigned i = 0 ; i < _offset.size() ; ++i )
      cout << _mitp->_sitp[i+_onset.size() ]->interpolant() <<" ";
   cout<< endl;
}
